use std::fs::read_to_string;

use crate::Puzzle;

#[derive(PartialEq, Debug)]
pub struct TheTreacheryOfWhales {
    positions: Vec<u32>,
}

impl TheTreacheryOfWhales {
    fn from(raw_data: &str) -> Self {
        let positions = raw_data
            .split(",")
            .map(|p| p.parse::<u32>().unwrap())
            .collect::<Vec<u32>>();

        Self { positions }
    }
}

impl Puzzle for TheTreacheryOfWhales {
    type Answer = u32;

    fn new() -> Self {
        let raw_data = read_to_string("input/day07.md").unwrap();

        Self::from(&raw_data)
    }

    fn solve_part1(&self) -> Self::Answer {
        let mut sorted_positions = self.positions.clone();
        sorted_positions.sort();
        let middle_index = sorted_positions.len() / 2;
        let middle_value = sorted_positions[middle_index];

        sorted_positions.iter().fold(0_u32, |sum, position| {
            sum + ((*position as i32) - (middle_value as i32)).unsigned_abs()
        })
    }

    fn solve_part2(&self) -> Self::Answer {
        let sum: f64 = self.positions.iter().sum::<u32>().into();
        let length: f64 = (self.positions.len() as u32).into();
        let average = sum / length;

        // Try both floor and ceil of the average to find the minimum cost
        let target1 = average.floor() as u32;
        let target2 = average.ceil() as u32;

        let cost1: u32 = self
            .positions
            .iter()
            .map(|position| {
                let distance = ((*position as i32) - (target1 as i32)).unsigned_abs();
                distance * (distance + 1) / 2
            })
            .sum();

        let cost2: u32 = self
            .positions
            .iter()
            .map(|position| {
                let distance = ((*position as i32) - (target2 as i32)).unsigned_abs();
                distance * (distance + 1) / 2
            })
            .sum();

        cost1.min(cost2)
    }
}

#[cfg(test)]
mod test {
    use super::*;

    fn create() -> TheTreacheryOfWhales {
        TheTreacheryOfWhales::from("16,1,2,0,4,2,7,1,2,14")
    }

    #[test]
    fn from() {
        assert_eq!(
            create(),
            TheTreacheryOfWhales {
                positions: vec![16, 1, 2, 0, 4, 2, 7, 1, 2, 14],
            }
        );
    }

    #[test]
    fn part1() {
        assert_eq!(create().solve_part1(), 37)
    }

    #[test]
    fn part2() {
        assert_eq!(create().solve_part2(), 168)
    }
}
